{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HMinCaml.Inline
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Inline expansion.
--
--------------------------------------------------------------------------------

module HMinCaml.Inline
  (

    inline
  , ksize

  ) where

import HMinCaml.Alpha
import qualified HMinCaml.HMonad as H
import HMinCaml.Id
import HMinCaml.KNormal
import qualified HMinCaml.Type as T


import Control.Applicative
import Control.Monad
import qualified Data.Map as M


ksize :: Expr -> Int
ksize (IfEq _ _ e1 e2) = 1 + ksize e1 + ksize e2
ksize (IfLE _ _ e1 e2) = 1 + ksize e1 + ksize e2
ksize (Let _ e1 e2)    = 1 + ksize e1 + ksize e2
ksize (LetRec fd e2)   = 1 + ksize (fun_body fd) + ksize e2
ksize (LetTuple _ _ e) = 1 + ksize e
ksize _                = 1



--------------------------------------------------------------------------------
-- Inline expansion

type FunDefn = ([(IdS,T.Type)], Expr)

type Env = M.Map IdS FunDefn

newtype Mon a = Mon { getMon :: Env -> H.HMonad a }


instance Functor Mon where
  fmap f ma = Mon $ \r -> fmap f $ getMon ma r

instance Applicative Mon where
  pure a    = Mon $ \_ -> pure a
  af <*> av = Mon $ \r -> getMon af r <*> getMon av r

instance Monad Mon where
  return    = pure
  ma >>= k  = Mon $ \r -> getMon ma r >>= \a -> getMon (k a) r

instance MonadPlus Mon where
  mzero         = Mon $ \_ -> mzero
  ma `mplus` mb = Mon $ \r -> getMon ma r `mplus` getMon mb r

-- Not modular - TODO should not have a run function...
--
runMon :: Mon a -> H.HMonad a
runMon ma = getMon ma M.empty

local :: (Env -> Env) -> Mon a -> Mon a
local f ma = Mon $ \r -> getMon ma (f r)


localBinding :: IdS -> FunDefn -> Mon a -> Mon a
localBinding v body ma = local (M.insert v body) ma

findMb :: IdS -> Mon (Maybe FunDefn)
findMb x = Mon $ \r -> return (M.lookup x r)


inline :: Expr -> H.HMonad Expr
inline expr = runMon $ ie expr


ie :: Expr -> Mon Expr
ie (IfEq x y e1 e2)         = IfEq x y <$> ie e1 <*> ie e2
ie (IfLE x y e1 e2)         = IfLE x y <$> ie e1 <*> ie e2

ie (LetRec (Fundef (x,t) yts e1) e2) = 
    aboveThreshold (ksize e1) >>= \a -> 
    if a then localBinding x (yts,e1) rest else rest
  where
    rest = (\ea eb -> LetRec (Fundef (x,t) yts ea) eb) <$> ie e1 <*> ie e2

ie (App x ys)               = findMb x >>= \ans -> case ans of
    Nothing     -> return (App x ys)
    Just (zs,e) -> let env = foldr fn M.empty (zip zs ys)
                   in Mon $ \_ -> alpha' env e
  where
    fn ((z,_),y) ac = M.insert z y ac


ie (LetTuple xts y e)       = LetTuple xts y <$> ie e
ie e                        = pure e


aboveThreshold :: Int -> Mon Bool
aboveThreshold i = Mon $ \_ -> H.threshold >>= \j -> return (i > j)


